package com.stars.pao.aicloud.controller;

import com.alibaba.cloud.ai.tongyi.chat.TongYiChatClient;
import com.alibaba.cloud.ai.tongyi.image.TongYiImagesClient;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.stars.pao.aicloud.entity.Answer;
import com.stars.pao.aicloud.entity.enums.AiModelEnum;
import com.stars.pao.aicloud.entity.enums.AiTypeEnum;
import com.stars.pao.aicloud.service.IAnswerService;
import com.stars.pao.aicloud.util.AppVariable;
import com.stars.pao.aicloud.util.MinIoUtil;
import com.stars.pao.aicloud.util.ResponseEntity;
import com.stars.pao.aicloud.util.SecurityUtil;
import com.stars.pao.aicloud.util.idempotent.Idempotent;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.Image;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.util.Base64;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

@Tag(name = "通义大模型控制器")
//通义大模型调用
@RestController
@RequestMapping("/tongyi")
public class TongyiController {

    @Resource
    private TongYiChatClient chatClient;
    @Resource
    private TongYiImagesClient imagesClient;
    @Resource
    private IAnswerService  answerService;
    @Resource
    private MinIoUtil minIoUtil;
    @Resource
    private RedissonClient redissonClient;
    @Resource
    private RedisTemplate<String, Object>  redisTemplate;

    /**
     * 聊天
     *
     * @param question
     * @return
     */
    @Idempotent
    @PostMapping("/chat")
    public ResponseEntity chat(String question) {
        if (!StringUtils.hasLength(question)) {
            return ResponseEntity.fail("问题不能为空");
        }
        // 执行分布式锁
        Long uid = SecurityUtil.getCurrentUser().getUid();
        String lockKey = AppVariable.getModelLockKey(uid,
                AiModelEnum.TONGYI.getValue(),
                AiTypeEnum.CHAT.getValue());
        String result= ""; // 大模型返回的结构
        boolean isSave = false; // 数据添加状态
        // 获取分布式锁实例
        RLock rLock = redissonClient.getLock(lockKey);
        try {
            boolean isLock = rLock.tryLock(30, TimeUnit.SECONDS);
            if (!isLock) {
                // 获取分布式锁失败
                return ResponseEntity.fail("请勿频繁请求，请稍后再试");
            }
            result = chatClient.call(new Prompt(question))
                    .getResult()
                    .getOutput()
                    .getContent();
            //        保存数据在数据库中
            Answer answer = new Answer();
            answer.setTitle(question);
            answer.setContent(result);
            answer.setModel(AiModelEnum.TONGYI.getValue());
            answer.setType(AiTypeEnum.CHAT.getValue());
            answer.setUid(uid);
            isSave = answerService.save(answer);
        } catch (Exception e) {
        } finally {
            rLock.unlock();
        }
        if(!isSave) {
//            保存失败
            return ResponseEntity.fail("数据保存失败，请重试");
        }
        return ResponseEntity.succ(result);
    }

    @Idempotent
    @PostMapping("/draw")
    public ResponseEntity draw(String question) {
        if (!StringUtils.hasLength(question)) {
            return ResponseEntity.fail("问题不能为空");
        }
//        String result = imagesClient.call(new ImagePrompt(question))
//                .getResult()
//                .getOutput()
//                .getUrl();
        // 执行分布式锁
        Long uid = SecurityUtil.getCurrentUser().getUid();
        String lockKey = AppVariable.getModelLockKey(uid,
                AiModelEnum.TONGYI.getValue(),
                AiTypeEnum.DRAW.getValue());
        String url= ""; // 大模型返回的结构
        boolean isSave = false; // 数据添加状态
        // 获取分布式锁实例
        RLock rLock = redissonClient.getLock(lockKey);
        try {
            boolean isLock = rLock.tryLock(30, TimeUnit.SECONDS);
            if (!isLock) {
                // 获取分布式锁失败
                return ResponseEntity.fail("请勿频繁请求，请稍后再试");
            }
            // 不获取url会过期
            Image image = imagesClient.call(new ImagePrompt(question))
                    .getResult()
                    .getOutput();
            String base64 = image.getB64Json();
            String fileName = "ty-" + UUID.randomUUID().toString().replace("-", "");
            try (InputStream inputStream = new ByteArrayInputStream(Base64.getDecoder().decode(base64))) {
                url = minIoUtil.upload(fileName, inputStream, "image/png");
            } catch (Exception e) {
                return ResponseEntity.fail("图片保存失败，请重试");
            }
            //        保存数据在数据库中
            Answer answer = new Answer();
            answer.setTitle(question);
            answer.setContent(url);
            answer.setModel(AiModelEnum.TONGYI.getValue());
            answer.setType(AiTypeEnum.DRAW.getValue());
            answer.setUid(uid);
            isSave = answerService.save(answer);
        } catch (Exception e) {
        } finally {
            rLock.unlock();
        }
        if(!isSave) {
            // 保存失败
            return ResponseEntity.fail("数据保存失败，请重试");
        }
        return ResponseEntity.succ(url);
    }

    //获取聊天历史列表
    @PostMapping("/getChatList")
    public ResponseEntity getChatList() {
        Long uid = SecurityUtil.getCurrentUser().getUid();
        int model = AiModelEnum.TONGYI.getValue();
        int type = AiTypeEnum.CHAT.getValue();
        String cacheKey = AppVariable.getListCacheKey(uid, model,type);
        Object list = redisTemplate.opsForValue().get(cacheKey);
        if(list == null)  {
            // 缓存不在
            QueryWrapper<Answer> queryWrapper = new QueryWrapper<>();
            queryWrapper.eq("uid", uid);
            queryWrapper.eq("model", model);
            queryWrapper.eq("type", type);
            queryWrapper.orderByDesc("aid");
            List<Answer> datalist = answerService.list(queryWrapper);
            redisTemplate.opsForValue().set(cacheKey, datalist, 1, TimeUnit.DAYS);
            return ResponseEntity.succ(datalist);
        } else {
            return ResponseEntity.succ(list);
        }
    }

    //获取绘画历史列表
    @PostMapping("/getDrawList")
    public ResponseEntity getDrawList() {
        Long uid = SecurityUtil.getCurrentUser().getUid();
        int model = AiModelEnum.TONGYI.getValue();
        int type = AiTypeEnum.DRAW.getValue();
        String cacheKey = AppVariable.getListCacheKey(uid, model,type);
        Object list = redisTemplate.opsForValue().get(cacheKey);
        if(list == null) {
            QueryWrapper<Answer> queryWrapper = new QueryWrapper<>();
            queryWrapper.eq("uid", SecurityUtil.getCurrentUser().getUid());
            queryWrapper.eq("model",model);
            queryWrapper.eq("type", type);
            queryWrapper.orderByDesc("aid");
            List<Answer> dataList = answerService.list(queryWrapper);
            redisTemplate.opsForValue().set(cacheKey, dataList, 1, TimeUnit.DAYS);
            return ResponseEntity.succ(dataList);
        } else {
            return ResponseEntity.succ(list);
        }
    }
}

